package com.polaris.live.common.util

import com.google.common.cache.CacheBuilder
import kotlinx.coroutines.sync.Mutex
import kotlinx.coroutines.sync.withLock
import java.util.concurrent.ConcurrentHashMap

/**
 * MutexUtils
 *
 * @author Created by lucas on 2024/5/30 07:43
 */

/**
 * 互斥的Map
 *
 * @author Created by lucas on 2024/5/30 07:44
 * @since 1.0.0
 */
class MutexMap<K : Any, V> {

    private val map = ConcurrentHashMap<K, V>()
    private val mutexes = Mutexes<K>()

    /**
     * 获取缓存
     *
     * 如果为空不会缓存
     */
    suspend fun get(key: K, block: suspend (K) -> V?): V? {
        return map[key] ?: mutexes.withLock(key) {
            map[key] ?: block(key)?.apply {
                map[key] = this
            }
        }
    }
}

/**
 * 多个互斥锁
 *
 * @author Created by lucas on 2024/5/30 07:40
 * @since 1.0.0
 */
class Mutexes<K : Any> {

    private val mutexes = ConcurrentHashMap<K, Mutex>()

    /**
     * get mutex by key
     */
    fun getMutex(key: K): Mutex = mutexes.computeIfAbsentBy21(key) { Mutex() }

    /**
     * Executes the given [action] under this mutex's lock.
     *
     * @param key the key to lock on.
     * @param owner Optional owner token for debugging. When `owner` is specified (non-null value) and this mutex
     *        is already locked with the same token (same identity), this function throws [IllegalStateException].
     *
     * @return the return value of the action.
     */
    suspend fun <T> withLock(key: K, owner: Any? = null, action: suspend () -> T): T {
        return getMutex(key).withLock(owner ?: key) {
            action()
        }
    }
}

/**
 * 互斥的Cache
 *
 * @author Created by lucas on 2024/5/30 07:44
 * @since 1.0.0
 */
class MutexCache<K : Any, V>(
    builder: () -> CacheBuilder<Any, Any>
) {

    private val cache = builder().build<K, V>()
    private val mutexes = builder().build<K, Mutex>()

    /**
     * 获取缓存
     *
     * 如果为空不会缓存
     */
    suspend fun get(key: K, block: suspend (K) -> V?): V? {
        return cache.getIfPresent(key) ?: withLock(key) {
            cache.getIfPresent(key) ?: block(key)?.apply {
                cache.put(key, this)
            }
        }
    }

    private fun getMutex(key: K): Mutex = mutexes.get(key) { Mutex() }

    private suspend fun <T> withLock(key: K, owner: Any? = null, action: suspend () -> T): T {
        return getMutex(key).withLock(owner ?: key) {
            action()
        }
    }
}